import aiohttp
import json
import asyncio
import logging
from typing import AsyncGenerator
from ..utils.http_client import HTTPClient
from core.config import settings

logger = logging.getLogger(__name__)


class OllamaService:
    """Ollama生成服务"""

    def __init__(self):
        self.client = HTTPClient()

    async def generate_stream(
            self,
            prompt: str,
            model: str = None
    ) -> AsyncGenerator[str, None]:
        """流式生成文本"""
        model = model or settings.OLLAMA_MODEL
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": True
        }

        session = await self.client.get_session()

        for attempt in range(settings.OLLAMA_MAX_RETRIES):
            try:
                async with session.post(
                        settings.OLLAMA_API_URL,
                        json=payload,
                        raise_for_status=True
                ) as response:
                    async for line in response.content:
                        line = line.decode('utf-8').strip()
                        if not line:
                            continue

                        try:
                            data = json.loads(line)
                            if "response" in data:
                                yield data["response"]
                            elif "error" in data:
                                logger.error(f"Ollama API error: {data['error']}")
                                yield f"Error: {data['error']}"
                        except json.JSONDecodeError:
                            logger.warning(f"Failed to parse Ollama response: {line}")
                            continue
                    return  # 成功完成，退出函数

            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
                if attempt < settings.OLLAMA_MAX_RETRIES - 1:
                    await asyncio.sleep(2 ** attempt)  # 指数退避
                else:
                    yield f"Request to Ollama API failed after {settings.OLLAMA_MAX_RETRIES} attempts: {str(e)}"
                    return